import sklearn
import argparse
import torch
from src.trainer import run_kfold_training
from src.analyzer import analyze_labels
from src.datasets import ClassificationDataset

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_path", type=str, default="/home/ma-user/work/Dataset/public_processed", help="原始数据主目录")
    parser.add_argument("--save_json_dir", type=str, default="./experiments/data_selector/", help="输出json文件路径")
    parser.add_argument("--kfold", type=int, default=4)
    parser.add_argument("--batch_size", type=int, default=100)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--wd", type=float, default=5e-4)
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--device", type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                        help="使用的设备，默认为cuda，如果没有可用的GPU则使用cpu")
    parser.add_argument("--model_cache_dir", type=str, default="experiments/huggingmodel")
    parser.add_argument("--output_dir", type=str, default="./experiments/outputs")
    parser.add_argument("--models", nargs='+', type=str, default=["efficientnet_b4", "resnet101"],
                        help="模型类型列表")
    args = parser.parse_args()

    # ==== 1. 加载数据集 ====
    print(f"📦 正在构建数据集 ...")
    dataset = ClassificationDataset(
        args.data_path,  args.save_json_dir
    )
    print(f"✅ 数据集加载完毕，样本数: {len(dataset)}")
    print(f"使用设备: {args.device}")
    # ==== 2. K折训练与推理 ====
    print(f"\n🎬 开始K折训练与推理，模型列表：{args.models}")
    probs, labels, path2prob, path2label = run_kfold_training(
        dataset=dataset,
        models=args.models,
        kfold=args.kfold,
        batch_size=args.batch_size,
        lr=args.lr,
        wd=args.wd,
        epochs=args.epochs,
        device=args.device,
        model_cache_dir=args.model_cache_dir
    )
    
    # ==== 3. 标签质量分析 ====
    print(f"\n🔬 开始标签质量分析 ...")
    analyze_labels(
        path2prob=path2prob,
        path2label=path2label,
        output_dir=args.output_dir
    )
    print("\n🎉 全部流程结束！请查看outputs和log下的json文件。")
if __name__ == '__main__':
    main()
